from sqlalchemy import Table, Column, Integer, String, MetaData, ForeignKey
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker


Base = declarative_base()

class Protein(Base):
    __tablename__ = 'protein'

    pid = Column(String, primary_key=True)
    start = Column(Integer)
    strand = Column(String)
    chromosome = Column(String)
    taxon = Column(String)

    def __init__(self, pid, start, strand, chromosome, taxon):
        self.pid = pid
        self.start = start
        self.strand = strand
        self.chromosome = chromosome
        self.taxon = taxon

    def __repr__(self):
        return "<Protein('%s')>" % (str(self.pid))


class OrthPair(Base):
    __tablename__ = 'orthpair'

    pairid = Column(Integer, primary_key=True)
    pid1 = Column(String, ForeignKey('protein.pid'), index=True)
    pid2 = Column(String, ForeignKey('protein.pid'))

    def __init__(self, pairid, pid1, pid2):
        self.pairid = pairid
        self.pid1 = pid1
        self.pid2 = pid2

    def __repr__(self):
        return "<OrthPair('%s','%s')>" % (self.pid1, self.pid2)


class DataBase(object):

    def __init__(self, protein_file, orth_file, engine='sqlite:///:memory:'):
        # need some checks here
        self.protein_file = protein_file
        self.orth_file = orth_file

        self.protein_table = Protein.__table__
        self.orthpair_table = OrthPair.__table__

        self.init(engine)

        self.dump_protein(self.protein_file)
        self.dump_orthpair(self.orth_file)


    def init(self, engine):
        self.engine = create_engine(engine, echo=False)
        self.metadata = Base.metadata
        self.metadata.create_all(self.engine)
        self.conn = self.engine.connect()
        Session = sessionmaker(bind=self.engine)
        self.session = Session()

    def dump_protein(self, filename):
        #need some check

        count = 0
        tmp = list()
        for line in open(filename):
            pid, start, end, strand, chromosome, taxon = line.strip().split('\t')
            if start == '':
                continue

            tmp.append({'pid':str(pid), 'start':int(start), 'strand':str(strand),
                        'chromosome':str(chromosome), 'taxon':str(taxon)})
            count += 1
            if count % 10000 == 0:
                self.conn.execute(self.protein_table.insert(), tmp)
                tmp = list()
        if tmp:
            self.conn.execute(self.protein_table.insert(), tmp)


    def dump_orthpair(self, filename):
        #also need some checks

        count = 1
        tmp = list()
        for line in open(filename):
            pid1, pid2 = line.strip().split('\t')[0:2]
            tmp.append({'pairid':int(count), 'pid1':str(pid1), 'pid2':str(pid2)})
            count += 1
            tmp.append({'pairid':int(count), 'pid1':str(pid2), 'pid2':str(pid1)})
            count += 1
            if count % 10000 == 0:
                self.conn.execute(self.orthpair_table.insert(), tmp)
                tmp = list()

        if tmp:
            self.conn.execute(self.orthpair_table.insert(), tmp)


    def get_profile(self, taxon):
        profile = dict()
        for pid, in self.session.query(Protein.pid).filter(Protein.taxon==taxon):
            profile[pid] = list()
            for orth_pid, in self.session.query(OrthPair.pid2).filter(OrthPair.pid1==pid):
                for t, in self.session.query(Protein.taxon).filter(Protein.pid==orth_pid):
                    profile[pid].append(t)

        return profile




if __name__ == '__main__':
    #this is only test
    import os
    os.chdir('..')
    os.chdir('Test')
    protein_file = 'protein_sample'
    orth_file = 'orth_sample'

    test_db = DataBase(protein_file, orth_file)
    test_db.get_profile('284590')
    raw_input()




#END